#!/usr/bin/env python3
# coding=utf-8
import rospy
from std_msgs.msg import String
import socket, struct, time, threading

# 从原始代码中导入类
from DogMotion import HeartBeat,ControlRobot

class ControlRobotROS:
    """
    ROS节点，用于通过话题控制机器狗并发布其状态。
    """
    def __init__(self, robot_ip='192.168.1.120', robot_port=43893):
        # 初始化ROS节点
        rospy.init_node('dog_controller_node', anonymous=True)
        rospy.loginfo("机器人控制节点启动...")
        
        # 实例化原始机器人控制类
        # 这个类包含了心跳、状态接收和命令发送的逻辑
        self.control = ControlRobot(ip=robot_ip, port=robot_port)
        self.control.autonomous_mode()
        rospy.loginfo("已切换到自主模式...")
        ######################
        # 留出检查机器狗是否起立（未完成）
        ######################

        # 订阅YOLO识别结果话题，消息形如 "color-number-state"
        # 提取state字段，根据state决定执行点头或摇头
        self.yolo_result_subscriber = rospy.Subscriber(
            '/yolo_result', 
            String, 
            self.yolo_result_callback
        )
        rospy.loginfo("已订阅 /yolo_result 话题...")
        # 发布机器人状态话题
        # 发布String类型消息，消息名为 '/dog_action_state'
        self.state_publisher = rospy.Publisher(
            '/dog_action_state', 
            String, 
            queue_size=10
        )
        rospy.loginfo("已创建 /dog_action_state 话题发布器...")
        # 注册节点关闭回调
        rospy.on_shutdown(self.on_shutdown)
    def nod(self):
        self.control.on_spot_mode()
        rospy.loginfo("已切换到原地模式...")
        self.control.move_forward(15000)
        self.control.move_forward(15000)
        rospy.loginfo("已执行点头...")
        self.control.move_mode()    
    def shake(self):
        self.control.on_spot_mode()
        rospy.loginfo("已切换到原地模式...")
        self.control.move_side(25000)
        self.control.move_side(-25000)
        rospy.loginfo("已执行摇头...")
        self.control.move_mode()
    def pub_action_state(self, state):
        """
        发布机器人的当前状态到 '/dog_action_state' 话题。
        """
        msg = String()
        msg.data = state+'-ok'
        self.state_publisher.publish(msg)
        rospy.loginfo(f"已发布状态: {state}")
    def on_shutdown(self):
        """
        当节点被关闭（例如按下 Ctrl+C）时执行。
        在此处添加需要在退出时运行的代码。
        """
        self.control.manual_mode()
    def yolo_result_callback(self, msg):
        """
        处理YOLO识别结果，字符串格式: "color-number-state"。
        提取state字段，如果为"up"或"down"则点头，否则摇头。
        """
        raw = msg.data.strip()
        rospy.loginfo(f"收到 /yolo_result: {raw}")

        parts = raw.split('-') if raw else []
        state = parts[-1].lower() if len(parts) >= 3 else ''

        if state in ('up', 'down'):
            self.nod()
            self.pub_action_state('nod')
            
        else:
            self.shake()
            self.pub_action_state('shake')
    

          
if __name__ == '__main__':
    try:
        # 创建ROS控制节点实例
        ros_controller = ControlRobotROS()
        # 保持节点运行，直到接收到关闭信号
        rospy.spin()
    except rospy.ROSInterruptException:
        rospy.loginfo("ROS节点已关闭。")
    except NameError:
        print("请确保robot_control_module.py文件存在，且包含HeartBeat, RobotState, ControlRobot类。")